import cv2
import numpy as np

class FourCorners:
    def __init__(self):
        self.left_top = (0,0) #左上角坐标
        self.left_bottom = (0,0) #左下角坐标
        self.right_top = (0,0) #右上角坐标
        self.right_bottom = (0,0) #右下角坐标
        
# 单应性变换，输入单应性变换矩阵和源图像形状，输出变换后的的角的坐标
def calc_corners(homography, src_shape):
    """
    计算源图像的四个角点经过单应性变换后的位置
    :param homography: 3x3 单应性矩阵
    :param src_shape: 源图像的形状(height, width)
    :return: FourCorners对象，也就是变换后四个角的坐标
    """
    h, w = src_shape # 高和宽
    corners = np.array([[0,0], [0,h-1], [w-1,0], [w-1,h-1]], dtype=np.float32) # 四个角做成数组，坐标从0开始，
    # 很烦，顺序又变成宽和高
    
    # 对角点应用透视变换
    # reshape(1,-1,2) 将角点数组变为 1xNx2 的形状以满足 perspectiveTransform 的输入要求
    # reshape(-1,2) 将变换后的结果变回 Nx2 的形状 (N个点，每个点2个坐标)
    transformed = cv2.perspectiveTransform(corners.reshape(1,-1,2), homography).reshape(-1,2)
    
    # 返回一个对象，包含四个点的坐标
    four_corners = FourCorners()
    four_corners.left_top = tuple(transformed[0])      # 左上角变换后坐标
    four_corners.left_bottom = tuple(transformed[1])   # 左下角变换后坐标
    four_corners.right_top = tuple(transformed[2])     # 右上角变换后坐标
    four_corners.right_bottom = tuple(transformed[3])  # 右下角变换后坐标
    return four_corners
    

# 接缝优化：重叠区域上做线性加权融合
def optimize_seam(img1, trans, img2, corners):
    """
    重叠区域上做线性加权融合
    :param img1: 源图像
    :param trans: 经过单应性变换后的图像
    :param img2: 目标图像
    :param corners: FourCorners对象，也就是变换后四个角的坐标
    """
    # 估算重叠区域的起始横坐标（取变换后 img1 左边界的最小 x 值）
    start = min(corners.left_top[0], corners.left_bottom[0])
    width = img2.shape[1] - int(start)
    
    rows, cols = img2.shape[:2]
    alpha = 1
    
    for row in range(rows):
        for col in range(cols):
            # 仅在重叠区域进行融合
            if col < start:
                continue
                
            # 两图的对应像素
            p = img2[row, col]
            t = trans[row, col]
            
            # 如果变换后的像素为黑色，直接使用原图
            if np.all(t == 0):
                alpha = 1
            else:
                alpha = (width - (col - start)) / width
                
            # 像素加权融合
            img2[row, col] = p * alpha + t * (1 - alpha)
    
    return img2

def stitch_images(img1, img2):
    """
    将两张图像 img1 (源) 和 img2 (目标) 拼接在一起。
    :param img1: 源图像 (通常是左图)
    :param img2: 目标图像 (通常是右图)
    :return: 拼接后的全景图像
    """
    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    
    # 使用ORB算法检测关键点和计算描述子
    orb = cv2.ORB_create(3000) # 指定期望的最大关键点数量
    print('开始计算特征点(keypoints)和描述子(descriptors)...')
    kp1, des1 = orb.detectAndCompute(gray1, None)
    kp2, des2 = orb.detectAndCompute(gray2, None) # 这里具体的数据结构是什么？
    print(f'img1的特征点和描述子数量为: {len(kp1)} 和 {des1.shape if des1 is not None else "None"}')
    print(f'img2的特征点和描述子数量为: {len(kp2)} 和 {des2.shape if des2 is not None else "None"}')
    
    # 特征匹配，使用FLANN匹配器
    # 配置FLANN参数，这里使用的是LSH (Locality Sensitive Hashing) 索引，它特别适合于匹配ORB等二进制描述子。index_params 定义了LSH索引的构建方式
    print('开始使用FLANN匹配器进行特征匹配...')
    FLANN_INDEX_LSH = 6
    index_params = dict(algorithm=FLANN_INDEX_LSH,
                        table_numbers=6,     # 哈希表数量
                        key_size=12,         # 哈希键大小
                        multi_probe_level=1) # 探针级别
    search_params = dict(checks=50)          # 搜索检查次数，越高越精确但是越慢
    
    # 创建FLANN匹配器
    matcher = cv2.FlannBasedMatcher(index_params, search_params)
    # 对 img1 中的每个描述子 des1，在 des2 中查找 k=2 个最相似的描述子。返回一个包含匹配对列表的列表。
    matches = matcher.knnMatch(des1, des2, k=2)
    print(f'匹配对数量为: {len(matches)}')
    # 扔掉不合格的匹配点，筛选
    good_matches = []
    for match_pair in matches:
        if len(match_pair) < 2:
            continue
        m, n = match_pair
        # 如果最佳匹配距离显著小于次佳匹配距离 (比率 < 0.5)，则认为是好的匹配
        if(m.distance < 0.5 * n.distance):
            good_matches.append(m)
    # 优质匹配点太少，无法进行准确拼接
    if len(good_matches) < 10:
        print('错误：找到的优质匹配点不足')
        return None
    print(f'找到的优质匹配点数量为: {len(good_matches)}')
    
    
    # 提取匹配点坐标
    # queryIdx 对应 des1 (img1) 中的索引, trainIdx 对应 des2 (img2) 中的索引
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1,1,2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1,1,2)
    
    # 计算单应性矩阵H，使用RANSAC算法剔除外点，H将点从img1坐标系映射到img2坐标系
    # H 是一个 3x3 的矩阵，将点从 img1 的坐标系映射到 img2 的坐标系
    print('开始计算单应性矩阵H...')
    H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0) # 5.0 是 RANSAC 的重投影误差阈值
    
    # 计算变换后img1的角点，用于确定画布大小
    corners = calc_corners(H, img1.shape[:2])
    print(f'经过H变换后的img1的角点为: {corners.left_top}, {corners.left_bottom}, {corners.right_top}, {corners.right_bottom}')
    
    # 计算拼接画布大小（添加强制类型转换）
    x_min = min(corners.left_top[0], corners.left_bottom[0], 0)
    x_max = max(corners.right_top[0], corners.right_bottom[0], img2.shape[1])
    y_min = min(corners.left_top[1], corners.right_top[1], 0)
    y_max = max(corners.left_bottom[1], corners.right_bottom[1], img2.shape[0])

    # 强制转换为整数
    x_min = int(round(x_min))
    x_max = int(round(x_max))
    y_min = int(round(y_min))
    y_max = int(round(y_max))
    # 得到最终的画布的宽和高
    width = x_max - x_min
    height = y_max - y_min
    print(f'最终画布的宽和高为: {width} 和 {height}')
    
    # 变换矩阵，将整个坐标系移动，使得左上角为 (0,0)
    translation_matrix = np.array([[1, 0, -x_min],  # 平移 x
                                   [0, 1, -y_min],  # 平移 y
                                   [0, 0, 1]])
    # img1会先经过H变换，再经过上面这个变换
    final_matrix = translation_matrix.dot(H)
    
    # 应用透视变换，将img1扭曲到最终画布坐标系
    warped = cv2.warpPerspective(img1, final_matrix, (width, height))
    
    # 创建最终空白画布
    result = np.zeros((height, width, 3), dtype=np.uint8)
    
    # 调整后的切片操作（添加显式类型转换）
    y_start = int(-y_min)
    y_end = y_start + img2.shape[0]
    x_start = int(-x_min)
    x_end = x_start + img2.shape[1]
    
    # 限制索引不超过数组边界
    y_start = max(0, y_start)
    y_end = min(result.shape[0], y_end)
    x_start = max(0, x_start)
    x_end = min(result.shape[1], x_end)
    
    # 计算对应的源图像区域
    src_y_start = max(0, -y_start)
    src_y_end = img2.shape[0] - max(0, y_end - result.shape[0])
    src_x_start = max(0, -x_start)
    src_x_end = img2.shape[1] - max(0, x_end - result.shape[1])
    
    result[y_start:y_end, x_start:x_end] = img2[src_y_start:src_y_end, src_x_start:src_x_end]
    
    # 优化接缝
    mask_warped = (warped != 0).astype(np.uint8)
    overlap = cv2.bitwise_and(mask_warped, (result != 0).astype(np.uint8))
    
    # 融合重叠区域
    fused = optimize_seam(warped, result, result.copy(), corners)
    
    # 拼接非重叠区域
    result[warped != 0] = warped[warped != 0]
    result[overlap == 1] = fused[overlap == 1]
    
    return result